import sys
import numpy as np
from PyQt5.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QFrame, QMainWindow
)
from PyQt5.QtChart import QChart, QChartView, QLineSeries, QValueAxis
from PyQt5.QtGui import QPen, QColor, QPainter, QDrag
from PyQt5.QtCore import Qt, QPointF, QMimeData

# 模拟数据
element_to_simulate = {
    "Mg": (0, 15, QColor(255,0,0)),
    "Al": (5, 20, QColor(0,150,255)),
    "Si": (10, 25, QColor(0,200,0)),
}
num_points = 50
well_id = 'A1'
depths = np.linspace(1000, 2000, num_points)
data = []
for i in range(num_points):
    row = {'井号': well_id, '深度': depths[i]}
    for element, (min_val, max_val, _) in element_to_simulate.items():
        row[element] = np.random.uniform(min_val, max_val)
    data.append(row)

class CurveLabel(QLabel):
    def __init__(self, name, color, parent=None):
        super().__init__(name, parent)
        self.setFixedHeight(24)
        self.setStyleSheet(f"background-color: {color.name()}; color: white; border-radius: 4px; padding:2px 8px;")
        self.setCursor(Qt.CursorShape.OpenHandCursor)
        self.name = name
        self.color = color
        self.drag_start_pos = None

    def mousePressEvent(self, a0):
        if a0 is not None and hasattr(a0, 'button') and a0.button() == Qt.MouseButton.LeftButton:
            self.drag_start_pos = a0.pos()
        super().mousePressEvent(a0)

    def mouseMoveEvent(self, a0):
        if a0 is not None and self.drag_start_pos is not None and hasattr(a0, 'pos') and (a0.pos() - self.drag_start_pos).manhattanLength() > 10:
            drag = QDrag(self)
            mime = QMimeData()
            mime.setText(self.name)
            drag.setMimeData(mime)
            drag.setProperty('curve_label', self)
            self.hide()
            drag.exec_(Qt.DropAction.MoveAction)
            self.show()
            self.drag_start_pos = None
        if a0 is not None:
            super().mouseMoveEvent(a0)

class TrackWidget(QFrame):
    def __init__(self, series_info_list, xdata, parent=None):
        super().__init__(parent)
        self.setFrameShape(QFrame.StyledPanel)
        self.setAcceptDrops(True)
        self.series_info_list = series_info_list  # [(name, color, ydata)]
        self.xdata = xdata
        self.init_ui()

    def init_ui(self):
        layout = QHBoxLayout(self)
        # y轴标签区
        self.label_area = QWidget()
        self.label_layout = QVBoxLayout()
        self.label_layout.setSpacing(2)
        self.label_area.setLayout(self.label_layout)
        self.curve_labels = []
        for name, color, ydata in self.series_info_list:
            label = CurveLabel(name, color)
            self.label_layout.addWidget(label)
            self.curve_labels.append(label)
        layout.addWidget(self.label_area)
        # 折线图
        self.chart = QChart()
        self.chart.setBackgroundVisible(False)
        self.axisX = QValueAxis()
        self.axisX.setTitleText("")
        self.axisX.setLabelFormat("%.1f")
        self.axisX.setTickCount(6)
        self.axisX.setRange(min(self.xdata), max(self.xdata))
        self.axisY = QValueAxis()
        self.axisY.setTitleText("")
        self.axisY.setLabelFormat("%.1f")
        self.axisY.setTickCount(6)
        # y轴范围覆盖所有曲线
        all_y = []
        for _, _, ydata in self.series_info_list:
            all_y.extend(ydata)
        self.axisY.setRange(min(all_y), max(all_y))
        self.chart.addAxis(self.axisX, Qt.AlignmentFlag.AlignBottom)
        self.chart.addAxis(self.axisY, Qt.AlignmentFlag.AlignLeft)
        self.series_objs = []
        for name, color, ydata in self.series_info_list:
            series = QLineSeries()
            series.setName(name)
            pen = QPen(color)
            pen.setWidth(2)
            series.setPen(pen)
            for x, y in zip(self.xdata, ydata):
                series.append(QPointF(x, y))
            self.chart.addSeries(series)
            series.attachAxis(self.axisX)
            series.attachAxis(self.axisY)
            self.series_objs.append(series)
        self.chart_view = QChartView(self.chart)
        self.chart_view.setRenderHint(QPainter.Antialiasing)
        self.chart_view.setFixedHeight(120)
        layout.addWidget(self.chart_view)
        self.setLayout(layout)
        self.setMouseTracking(True)

    def add_curve(self, name, color, ydata):
        # 添加曲线到当前轨道
        self.series_info_list.append((name, color, ydata))
        label = CurveLabel(name, color)
        self.label_layout.addWidget(label)
        self.curve_labels.append(label)
        # 添加到图表
        series = QLineSeries()
        series.setName(name)
        pen = QPen(color)
        pen.setWidth(2)
        series.setPen(pen)
        for x, y in zip(self.xdata, ydata):
            series.append(QPointF(x, y))
        self.chart.addSeries(series)
        series.attachAxis(self.axisX)
        series.attachAxis(self.axisY)
        self.series_objs.append(series)
        # 更新y轴范围
        all_y = []
        for _, _, ydata in self.series_info_list:
            all_y.extend(ydata)
        self.axisY.setRange(min(all_y), max(all_y))

    def remove_curve(self, name):
        # 从当前轨道移除曲线
        idx = None
        for i, (n, _, _) in enumerate(self.series_info_list):
            if n == name:
                idx = i
                break
        if idx is not None:
            self.series_info_list.pop(idx)
            label = self.curve_labels.pop(idx)
            label.setParent(None)
            series = self.series_objs.pop(idx)
            self.chart.removeSeries(series)
            # 更新y轴范围
            all_y = []
            for _, _, ydata in self.series_info_list:
                all_y.extend(ydata)
            if all_y:
                self.axisY.setRange(min(all_y), max(all_y))

    def get_curve_info(self, name):
        for n, color, ydata in self.series_info_list:
            if n == name:
                return (n, color, ydata)
        return None

class MultiTrackWidget(QWidget):
    def __init__(self, tracks_data, xdata):
        super().__init__()
        self.xdata = xdata
        self.track_widgets = []
        self.drag_indicator = QFrame()
        self.drag_indicator.setFixedHeight(10)
        self.drag_indicator.setStyleSheet("background-color: yellow; border: 1px solid orange;")
        self.drag_indicator.hide()
        self.init_ui(tracks_data)

    def init_ui(self, tracks_data):
        self.vlayout = QVBoxLayout()
        self.vlayout.setSpacing(5)
        for name, ydata, color in tracks_data:
            track = TrackWidget([(name, color, ydata)], self.xdata)
            self.vlayout.addWidget(track)
            self.track_widgets.append(track)
        # 共用x轴标签
        self.xaxis_label = QLabel(" "*10 + "深度/时间轴" + " "*10)
        self.xaxis_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.xaxis_label.setStyleSheet("font-weight:bold;font-size:16px;")
        self.vlayout.addWidget(self.xaxis_label)
        self.setLayout(self.vlayout)
        self.setAcceptDrops(True)

    def dragEnterEvent(self, a0):
        if a0 is not None:
            a0.accept()

    def dragMoveEvent(self, a0):
        if a0 is None:
            return
        pos = a0.pos()
        insert_index, target_track = self._find_insert_index_and_target(pos)
        self._show_drag_indicator(insert_index)
        a0.accept()

    def dragLeaveEvent(self, a0):
        self.drag_indicator.hide()
        if a0 is not None and hasattr(a0, 'accept'):
            a0.accept()

    def dropEvent(self, a0):
        self.drag_indicator.hide()
        if a0 is None:
            return
        pos = a0.pos()
        insert_index, target_track = self._find_insert_index_and_target(pos)
        src = a0.source()
        # 拖拽曲线标签
        if isinstance(src, CurveLabel):
            # 找到原始轨道
            for track in self.track_widgets:
                if src in track.curve_labels:
                    name = src.name
                    info = track.get_curve_info(name)
                    if info is not None:
                        # 从原轨道移除
                        track.remove_curve(name)
                        # 如果原轨道没曲线了，移除轨道
                        if not track.series_info_list:
                            self.vlayout.removeWidget(track)
                            self.track_widgets.remove(track)
                        # 合并到目标轨道或新建轨道
                        if target_track is not None:
                            target_track.add_curve(*info)
                        else:
                            new_track = TrackWidget([info], self.xdata)
                            self.vlayout.insertWidget(insert_index, new_track)
                            self.track_widgets.insert(insert_index, new_track)
                    break
        # 拖拽整个轨道
        elif isinstance(src, TrackWidget):
            if target_track is not None and src != target_track:
                # 合并所有曲线到目标轨道
                for name, color, ydata in src.series_info_list[:]:
                    target_track.add_curve(name, color, ydata)
                    src.remove_curve(name)
                # 移除原轨道
                self.vlayout.removeWidget(src)
                self.track_widgets.remove(src)
            else:
                # 只是排序
                if src in self.track_widgets:
                    idx_from = self.vlayout.indexOf(src)
                    if idx_from < insert_index:
                        insert_index -= 1
                    self.vlayout.removeWidget(src)
                    self.vlayout.insertWidget(insert_index, src)
                    self.track_widgets.remove(src)
                    self.track_widgets.insert(insert_index, src)
        a0.accept()

    def _find_insert_index_and_target(self, pos):
        # 计算鼠标y坐标对应的插入位置和目标轨道
        for i, track in enumerate(self.track_widgets):
            if pos.y() < track.y() + track.height() // 2:
                # 判断是否在track区域内
                if track.y() < pos.y() < track.y() + track.height():
                    return i, track
                else:
                    return i, None
        return len(self.track_widgets), None

    def _show_drag_indicator(self, index):
        self.drag_indicator.hide()
        if self.vlayout.indexOf(self.drag_indicator) != -1:
            self.vlayout.removeWidget(self.drag_indicator)
        self.vlayout.insertWidget(index, self.drag_indicator)
        self.drag_indicator.show()

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("多轨道折线图重叠/独立拖拽示例")
        self.setGeometry(100, 100, 1000, 600)
        tracks_data = []
        for name, (vmin, vmax, color) in element_to_simulate.items():
            ydata = [row[name] for row in data]
            tracks_data.append((name, ydata, color))
        self.widget = MultiTrackWidget(tracks_data, [row['深度'] for row in data])
        self.setCentralWidget(self.widget)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    main_window = MainWindow()
    main_window.show()
    sys.exit(app.exec_())